#include <math.h>
#include <stdio.h>
#include <sys/stat.h>
#include <algorithm>
#include <cstdint>
#include <fstream>
#include <vector>
#include "face_post_process.h"
#include "typeConv.hpp"
#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

struct SCRFDConfig
{
  int num_anchor;
  std::vector<int> strides;
  // give class name back to user
};

static SCRFDConfig default_SCRFD_config = {
    2,
    {8, 16, 32},
};

static inline float half2float(int16_t ib)
{
  const int sig[2] = {1, -1};
  const float result = 5.96046e-08;

  int16_t s, e, m;
  s = (ib >> 15) & 0x1;
  e = (ib >> 10) & 0x1f;
  m = ib & 0x3ff;

  // added by puyang.wang@lynxi.com
  {
    if (0 == e)
      return sig[s] * m * result;
    else
    {
      union
      {
        unsigned int u32;
        float f32;
      } ou;

      e = (0x1f == e) ? 0xff : (e - 15 + 127);
      ou.u32 = (s << 31) | (e << 23) | (m << 13);
      return ou.f32;
    }
  }
}

static inline int16_t float2half(float value)
{
  int16_t ob;
  //            uint16 s,e,m;
  int16_t s, m;
  int e; // modified by puyang.wang@lynxi。修正接近于0的值不能正确转换的bug。
  // int16_t expBitNum = 5;
  // int16_t baseBitNum = 15 - expBitNum;
  int maxExp = 32;    // pow(2,expBitNum);
  int maxBase = 1024; // pow(2,baseBitNum);
  int biasExp = 15;   // maxExp/2 - 1;
  s = value < 0;
  double thd2;
  // thd1 = (maxBase-1)*1.0/maxBase * pow(2,(1 - biasExp));
  thd2 = 1.0 / maxBase * pow(2, (1 - biasExp));
  double x;
  bool inf_flag = 0;
  x = s ? -value : value;
  x = (x > 65504) ? 65504 : x;
  int16_t indA;
  indA = x < thd2 / 2;
  // indB = x > thd2/2;
  if (indA)
  {
    e = 0;
    m = 0;
    s = 0;
  }
  // if (indB)
  else // float为Nan转为half
  {
    union
    {
      float xl;
      unsigned int u32;
    } ou;
    ou.xl = log2(x);

    if (((ou.u32 >> 23) & 0xff) == 0xff) // float为inf转为half
    {
      e = maxExp - 1;
      if ((ou.u32 & 0x7fffff) == 0)
        inf_flag = 1;
      else
      {
        inf_flag = 0;
        s = (ou.u32 >> 31);
      }
    }
    else
    {
      e = biasExp + floor(ou.xl);
    }

    if (e > (maxExp - 1))
      printf("[double2uint16]Error: out of e range\n");
  }
  int16_t ind1, ind2;
  ind1 = e <= 0;
  ind2 = e > 0;
  if (ind1)
  {
    e = 0;
    m = round(x * pow(2, (biasExp - 1)) * maxBase);
  }
  if (ind2)
  {
    if (31 == e)
    {
      if (inf_flag)
        m = 0;
      else
        m = 1;
    }
    else
    {
      double xr;
      xr = x / pow(2, (e - biasExp)) - 1;
      m = round(xr * maxBase);
    }
  }

  ob = (s & 0x1) << 15 | (((e & 0x1f) << 10) + m);
  return ob;
}

static inline float intersection_area(const FaceDetectBox &a,
                                      const FaceDetectBox &b)
{
  if (a.xmin > b.xmax || a.xmax < b.xmin || a.ymin > b.ymax ||
      a.ymax < b.ymin)
  {
    // no intersection
    return 0.f;
  }

  float inter_width = std::min(a.xmax, b.xmax) - std::max(a.xmin, b.xmin);
  float inter_height = std::min(a.ymax, b.ymax) - std::max(a.ymin, b.ymin);

  return inter_width * inter_height;
}

static inline void qsort_descent_inplace(std::vector<FaceDetectBox> &datas,
                                         int left, int right)
{
  int i = left;
  int j = right;
  float p = datas[(left + right) / 2].score;

  while (i <= j)
  {
    while (datas[i].score > p)
      i++;

    while (datas[j].score < p)
      j--;

    if (i <= j)
    {
      // swap
      std::swap(datas[i], datas[j]);

      i++;
      j--;
    }
  }

  if (left < j)
    qsort_descent_inplace(datas, left, j);

  if (i < right)
    qsort_descent_inplace(datas, i, right);
}

static void qsort_descent_inplace(std::vector<FaceDetectBox> &datas)
{
  if (datas.empty())
    return;

  qsort_descent_inplace(datas, 0, static_cast<int>(datas.size() - 1));
}

static void nms_sorted_bboxes(const std::vector<FaceDetectBox> &bboxes,
                              std::vector<size_t> &picked,
                              float nms_threshold)
{
  picked.clear();

  const size_t n = bboxes.size();

  for (size_t i = 0; i < n; i++)
  {
    const FaceDetectBox &a = bboxes[i];

    int keep = 1;
    for (int j = 0; j < (int)picked.size(); j++)
    {
      const FaceDetectBox &b = bboxes[picked[j]];

      // intersection over union
      float inter_area = intersection_area(a, b);
      float union_area = a.area + b.area - inter_area;
      // LOG_PLUGIN_E("union_area: %.2f   %.2f\n", union_area, inter_area);
      //             float IoU = inter_area / union_area
      if (inter_area / union_area > nms_threshold)
        keep = 0;
    }

    if (keep)
      picked.push_back(i);
  }
}

static void generate_proposals(int feat_stride, const int16_t *score_blob,
                               const int16_t *bbox_blob,
                               const int16_t *landmark_blob,
                               int16_t prob_threshold,
                               std::vector<FaceDetectBox> &faceobjects, int w,
                               int h)
{
  for (int i = 0; i < h; i++)
  {
    for (int j = 0; j < w; j++)
    {
      for (int q = 0; q < default_SCRFD_config.num_anchor; q++)
      {
        int index = i * w * default_SCRFD_config.num_anchor +
                    j * default_SCRFD_config.num_anchor + q;
        const int16_t score = score_blob[index];

        if (score >= prob_threshold)
        {
          float cx = j * feat_stride;
          float cy = i * feat_stride;

          const int16_t *bbox = bbox_blob + 4 * index;
          const int16_t *landmark = landmark_blob + 10 * index;
          float dx = half2float(bbox[0]);
          float dy = half2float(bbox[1]);
          float dw = half2float(bbox[2]);
          float dh = half2float(bbox[3]);

          float x0 = cx - dx * feat_stride;
          float y0 = cy - dy * feat_stride;
          float x1 = cx + dw * feat_stride;
          float y1 = cy + dh * feat_stride;

          float prob = half2float(score);

          FaceDetectBox obj;
          obj.xmin = x0;
          obj.ymin = y0;
          obj.xmax = x1;
          obj.ymax = y1;
          obj.area = (y1 - y0) * (x1 - x0);
          obj.landmark[0] = cx + feat_stride * half2float(landmark[0]);
          obj.landmark[1] = cy + feat_stride * half2float(landmark[1]);
          obj.landmark[2] = cx + feat_stride * half2float(landmark[2]);
          obj.landmark[3] = cy + feat_stride * half2float(landmark[3]);
          obj.landmark[4] = cx + feat_stride * half2float(landmark[4]);
          obj.landmark[5] = cy + feat_stride * half2float(landmark[5]);
          obj.landmark[6] = cx + feat_stride * half2float(landmark[6]);
          obj.landmark[7] = cy + feat_stride * half2float(landmark[7]);
          obj.landmark[8] = cx + feat_stride * half2float(landmark[8]);
          obj.landmark[9] = cy + feat_stride * half2float(landmark[9]);
          obj.score = prob;
          obj.id = 0;

          faceobjects.push_back(obj);
        }
      }
    }
  }
}

// [800, 1] box score
// [800, 4] box xywh
// [800, 10] landmark pred
void decode_tensor(std::vector<FaceDetectBox> &dets, const int feat_stride,
                   const int16_t *score_tensor, const int16_t *bbox_tensor,
                   const int16_t *landmark_tensor, int16_t obj_thresh,
                   const FacePostProcessInfo_t *post_info, int grid_h,
                   int grid_w, const int layer_id)
{
  std::vector<FaceDetectBox> strideObjs;
  generate_proposals(feat_stride, score_tensor, bbox_tensor, landmark_tensor,
                     obj_thresh, strideObjs, grid_w, grid_h);

  dets.insert(dets.end(), strideObjs.begin(), strideObjs.end());
}

/**
 * @brief
 * 输出节点信息：
 * [12800, 1]   stride 8 bbox score
 * [3200,  1]   stride 16 bbox score
 * [800,   1]   stride 32 bbox score
 * [12800, 4]   stride 8 bbox xywh pred
 * [3200,  4]   stride 16 bbox xywh pred
 * [800,   4]   stride 32 bbox xywh pred
 * [12800, 10]  stride 8 landmark pred
 * [3200,  10]  stride 16 landmark pred
 * [800,   10]  stride 32 landmark pred
 * @param post_info
 * @return
 */

int lynFacePostProcess(FacePostProcessInfo_t *post_info)
{
  int16_t *output_tensor = nullptr;
  FaceDetectInfo *detectInfo = nullptr;
#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr)
  {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  detectInfo = (FaceDetectInfo *)lynPluginGetVirtAddr(post_info->detectInfo);
  if (detectInfo == nullptr)
  {
    LOG_PLUGIN_E("get detectInfo addr error\n");
    return -1;
  }
#else
  output_tensor = (int16_t *)post_info->output_tensor;
  detectInfo = post_info->detectInfo;
#endif

  int layer0_grid_h = post_info->height / default_SCRFD_config.strides[0];
  int layer0_grid_w = post_info->width / default_SCRFD_config.strides[0];

  int layer1_grid_h = post_info->height / default_SCRFD_config.strides[1];
  int layer1_grid_w = post_info->width / default_SCRFD_config.strides[1];

  int layer2_grid_h = post_info->height / default_SCRFD_config.strides[2];
  int layer2_grid_w = post_info->width / default_SCRFD_config.strides[2];

  int num_tensor_grid_size = layer0_grid_h * layer0_grid_w +
                             layer1_grid_h * layer1_grid_w +
                             layer2_grid_h * layer2_grid_w;

  int16_t *score_tensor_base = (int16_t *)output_tensor;
  int16_t *bbox_tensor_base =
      score_tensor_base +
      num_tensor_grid_size * default_SCRFD_config.num_anchor * 1;
  int16_t *landmark_tensor_base =
      bbox_tensor_base +
      num_tensor_grid_size * default_SCRFD_config.num_anchor * 4;

  int16_t *score_tensor0 = score_tensor_base;
  int16_t *bbox_tensor0 = bbox_tensor_base;
  int16_t *landmark_tensor0 = landmark_tensor_base;

  int16_t *score_tensor1 = score_tensor0 + layer0_grid_h * layer0_grid_w *
                                               default_SCRFD_config.num_anchor *
                                               1;
  int16_t *bbox_tensor1 = bbox_tensor0 + layer0_grid_h * layer0_grid_w *
                                             default_SCRFD_config.num_anchor *
                                             4;
  int16_t *landmark_tensor1 =
      landmark_tensor_base +
      layer0_grid_h * layer0_grid_w * default_SCRFD_config.num_anchor * 10;

  int16_t *score_tensor2 = score_tensor1 + layer1_grid_h * layer1_grid_w *
                                               default_SCRFD_config.num_anchor *
                                               1;
  int16_t *bbox_tensor2 = bbox_tensor1 + layer1_grid_h * layer1_grid_w *
                                             default_SCRFD_config.num_anchor *
                                             4;
  int16_t *landmark_tensor2 =
      landmark_tensor1 +
      layer1_grid_h * layer1_grid_w * default_SCRFD_config.num_anchor * 10;

  float resize_scale;
  if (post_info->ori_height > post_info->ori_width)
  {
    resize_scale = (float)post_info->height / post_info->ori_height;
  }
  else
  {
    resize_scale = (float)post_info->width / post_info->ori_width;
  }

  std::vector<FaceDetectBox> dets;

  int16_t box_score_threshold = float2half(post_info->score_threshold);

  // stride 8
  decode_tensor(dets, 8, score_tensor0, bbox_tensor0, landmark_tensor0,
                box_score_threshold, post_info, layer0_grid_h, layer0_grid_w,
                0);

  // stride 16
  decode_tensor(dets, 16, score_tensor1, bbox_tensor1, landmark_tensor1,
                box_score_threshold, post_info, layer1_grid_h, layer1_grid_w,
                1);

  // stride 32
  decode_tensor(dets, 32, score_tensor2, bbox_tensor2, landmark_tensor2,
                box_score_threshold, post_info, layer2_grid_h, layer2_grid_w,
                2);

  // global sort inplace
  qsort_descent_inplace(dets);

  // apply nms
  std::vector<size_t> picked;
  nms_sorted_bboxes(dets, picked, post_info->nms_threshold);

  detectInfo->boxesNum = std::min((int)picked.size(), BOX_MAX_NUM);

  for (int i = 0; i < detectInfo->boxesNum; i++)
  {
    int index = picked[i];

    // clip to image size
    float x0 = dets[index].xmin;
    float y0 = dets[index].ymin;
    float x1 = dets[index].xmax;
    float y1 = dets[index].ymax;

    x0 = std::max(std::min(x0, (float)post_info->width - 1), 0.f);
    y0 = std::max(std::min(y0, (float)post_info->height - 1), 0.f);
    x1 = std::max(std::min(x1, (float)post_info->width - 1), 0.f);
    y1 = std::max(std::min(y1, (float)post_info->height - 1), 0.f);

    detectInfo->boxes[i].xmin = x0 / resize_scale;
    detectInfo->boxes[i].ymin = y0 / resize_scale;
    detectInfo->boxes[i].xmax = x1 / resize_scale;
    detectInfo->boxes[i].ymax = y1 / resize_scale;

    if (static_cast<int>(detectInfo->boxes[i].xmax) > post_info->ori_width - 1)
      detectInfo->boxes[i].xmax = post_info->ori_width - 1;
    if (static_cast<int>(detectInfo->boxes[i].xmin) < 0)
      detectInfo->boxes[i].xmin = 0;
    if (static_cast<int>(detectInfo->boxes[i].ymax) > post_info->ori_height - 1)
      detectInfo->boxes[i].ymax = post_info->ori_height - 1;
    if (static_cast<int>(detectInfo->boxes[i].ymin) < 0)
      detectInfo->boxes[i].ymin = 0;

    int num = sizeof(dets[index].landmark) / sizeof(float) / 2;
    for (int j = 0; j < num; j++)
    {
      detectInfo->boxes[i].landmark[2 * j + 0] = dets[index].landmark[2 * j + 0] / resize_scale;
      detectInfo->boxes[i].landmark[2 * j + 1] = dets[index].landmark[2 * j + 1] / resize_scale;
    }

    detectInfo->boxes[i].score = dets[index].score;
    detectInfo->boxes[i].id = dets[index].id;
  }

  return 0;
}